import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from reflow.diffusers import EulerDummyScheduler

model_base = "checkpoints/SD-1-4"
reflow_model = "logs/online/sdv1-4_laion_1MPrompts_guidance7.5/checkpoints/score_model_s40000.pth"
lora_model = "logs/lora/lr1e-4_rank128_s10000"
device='cuda:0'

pipe = StableDiffusionPipeline.from_pretrained(model_base, torch_dtype=torch.float16,safety_checker=None,
        requires_safety_checker=False,)
pipe.unet.load_state_dict(torch.load(reflow_model, map_location='cpu'))
# pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.scheduler = EulerDummyScheduler()

pipe.unet.load_attn_procs(lora_model)
pipe.to(device)

inference_steps=50
guidance_scale=6
seed=2
prompt="a picture of a cartoon character with a sword"

generator=torch.Generator()
generator.manual_seed(seed)

for lora_scale in [0,0.25,0.5,0.75,1]:
    image = pipe(
        prompt, num_inference_steps=inference_steps, guidance_scale=guidance_scale, cross_attention_kwargs={"scale": lora_scale}
    ).images[0]
    image.save(f"tmp/scale{lora_scale}.jpg")
